import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


spr = 2
spc = 7
spn = 0
plt.figure(figsize=[12, 6])

pic_dir = '../../../../../large_data/pic/fruit/'
A = cv.imread(pic_dir + 'apple.jpeg')
B = cv.imread(pic_dir + 'orange.jpeg')

# generate Gaussian pyramid for A
sep('gpA')
G = A.copy()
gpA = [G]
for i in range(6):
    G = cv.pyrDown(G)
    gpA.append(G)
for i, G in enumerate(gpA):  # ori, 0, 1, 2, 3, 4, 5
    print(f'#{i}, {G.shape}')
    my_show_img(G, 'apple ' + str(i), lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

# generate Gaussian pyramid for B
sep('gpB')
G = B.copy()
gpB = [G]
for i in range(6):
    G = cv.pyrDown(G)
    gpB.append(G)
for i, G in enumerate(gpB):  # ori, 0, 1, 2, 3, 4, 5
    print(f'#{i}, {G.shape}')
    my_show_img(G, 'orange ' + str(i), lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

spr = 2
spc = 3
spn = 0
plt.figure(figsize=[12, 6])

# generate Laplacian Pyramid for A
lpA = [gpA[5]]
for i in range(5, 0, -1):
    GE = cv.pyrUp(gpA[i])
    print(gpA[i].shape, gpA[i - 1].shape, GE.shape)
    GE = cv.resize(GE, (gpA[i - 1].shape[0], gpA[i - 1].shape[1]))
    L = cv.subtract(gpA[i - 1], GE)
    lpA.append(L)
    my_show_img(L, 'apple ' + str(i), lambda x: cv.cvtColor(x, cv.COLOR_BGR2GRAY), cmap='gray')

spr = 2
spc = 3
spn = 0
plt.figure(figsize=[12, 6])

# generate Laplacian Pyramid for B
lpB = [gpB[5]]
for i in range(5, 0, -1):
    GE = cv.pyrUp(gpB[i])
    print(gpA[i].shape, gpA[i - 1].shape, GE.shape)
    GE = cv.resize(GE, (gpA[i - 1].shape[0], gpA[i - 1].shape[1]))
    L = cv.subtract(gpB[i - 1], GE)
    lpB.append(L)
    my_show_img(L, 'orange ' + str(i), lambda x: cv.cvtColor(x, cv.COLOR_BGR2GRAY), cmap='gray')

spr = 2
spc = 3
spn = 0
plt.figure(figsize=[12, 6])

# Now add left and right halves of images in each level
LS = []
i = -1
for la,lb in zip(lpA,lpB):
    i += 1
    rows,cols,dpt = la.shape
    ls = np.hstack((la[:,0:cols//2], lb[:,cols//2:]))
    LS.append(ls)
    my_show_img(ls, 'cat ' + str(i), lambda x: cv.cvtColor(x, cv.COLOR_BGR2GRAY), cmap='gray')

spr = 2
spc = 3
spn = 0
plt.figure(figsize=[12, 6])

# now reconstruct
ls_ = LS[0]
for i in range(1,6):
    ls_ = cv.pyrUp(ls_)
    ls_ = cv.add(ls_, LS[i])
    my_show_img(ls_, 'add ' + str(i), lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB), cmap='gray')

# image with direct connecting each half
real = np.hstack((A[:,:cols//2],B[:,cols//2:]))

cv.imshow('ls_', ls_)
cv.imshow('real', real)
